import logging
from typing import List, Any
from sqlalchemy import text, select, delete, bindparam
from sqlalchemy.dialects.postgresql import JSONB

from llama_index.core.schema import BaseNode
from llama_index.core.vector_stores import VectorStoreQuery
from llama_index.core.vector_stores.types import VectorStoreQueryResult
from llama_index.vector_stores.postgres import PGVectorStore

from axiom_boot.vector.store.vector_store_interface import VectorStoreInterface

logger = logging.getLogger(__name__)


class PGVectorStoreAdapter(VectorStoreInterface):
    def __init__(self, native_store: PGVectorStore, collection_name: str):
        self._native_store = native_store
        self._collection_name = collection_name

    async def add(self, nodes: List[BaseNode]) -> List[str]:
        return await self._native_store.async_add(nodes)

    async def delete(self, ref_doc_id: str, **kwargs: Any) -> bool:
        """
        根据 ref_doc_id (文档ID) 物理删除所有相关的节点。
        [最终修复] 回归并正确使用 llama-index 官方提供的 delete 方法。
        该方法在内部会正确处理基于 ref_doc_id 的删除，避免了我们手动操作 SQL 可能引入的
        类型转换错误和对内部属性的非法访问。
        """
        try:
            await self._native_store.adelete(ref_doc_id=ref_doc_id, **kwargs)
            logger.info(f"已成功触发对文档 '{ref_doc_id}' 相关节点的删除操作。")
            return True
            
        except Exception as e:
            logger.error(f"从集合 '{self._collection_name}' 删除文档 '{ref_doc_id}' 的节点时发生错误: {e}", exc_info=True)
            raise

    async def aget_nodes(self, node_ids: List[str], **kwargs: Any) -> List[BaseNode]:
        """Get nodes from vector store by delegating to the native implementation."""
        return await self._native_store.aget_nodes(node_ids=node_ids, **kwargs)

    async def adelete_nodes(
        self,
        node_ids: List[str],
        delete_from_docstore: bool = False,
        **delete_kwargs: Any,
    ) -> None:
        """Delete nodes from vector store."""
        async with self._native_store._async_session_maker() as session:
            stmt = delete(self._native_store.node_table).where(
                self._native_store.node_table.node_id.in_(node_ids)
            )
            await session.exec(stmt)
            await session.commit()

    async def query(self, query: VectorStoreQuery, **kwargs: Any) -> VectorStoreQueryResult:
        return await self._native_store.aquery(query, **kwargs)

    async def delete_collection(self):
        async with self._native_store._async_engine.connect() as connection:
            async with connection.begin():
                table_name = self._native_store.table_name
                drop_stmt = text(f'DROP TABLE IF EXISTS public."{table_name}" CASCADE;')
                await connection.execute(drop_stmt)
                logger.info(f"向量集合 (表) '{table_name}' 已被成功删除。")
    
    async def collection_exists(self) -> bool:
        async with self._native_store._async_engine.connect() as conn:
            return await conn.run_sync(
                self._native_store._engine.dialect.has_table,
                self._native_store.table_name,
            )

    @property
    def client(self) -> Any:
        return self._native_store.client